
from langchain.document_loaders import TextLoader,JSONLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from faissManager import FAISS
from sent2vec import Sent2VecEmbeddings
from langchain.document_transformers import (
    LongContextReorder,
)
import os
import torch
index_name="/data/lxy/RAT/10m9d_wxb_bge1.5_hnswivf_20500"
# Load the document, split it into chunks, embed each chunk and load it into the vector store.
maxNorm=torch.load(os.path.join(index_name,'maxNorm.pt'))
model=Sent2VecEmbeddings(model_name='BAAI/bge-large-zh-v1.5')
db=FAISS.load_local(index_name,model)

# with open('补充数据.txt','r',encoding='utf-8') as f:
#     data=f.readlines()

print('添加前的数据量：',db.index.ntotal)
# [[doc,query]]!!!!
data=[['业务名称：游泳馆预约\n\n办理流程：用于游泳馆预约\n\n办理条件：东北大学师生\n\n负责部门：体育场馆管理中心\n\n办理网址：https://ehall.neu.edu.cn//db_portal/guide?id=9AA0F0B8-B2DA-4723-9646-A461EFF1D920','我是校外人员，可以用游泳馆吗？'],
      ['业务名称：游泳馆预约\n\n办理流程：用于游泳馆预约\n\n办理条件：东北大学师生\n\n负责部门：体育场馆管理中心\n\n办理网址：https://ehall.neu.edu.cn//db_portal/guide?id=9AA0F0B8-B2DA-4723-9646-A461EFF1D920','东北大学游泳馆需要预约嘛'],
      ['业务名称：游泳馆预约\n\n办理流程：用于游泳馆预约\n\n办理条件：东北大学师生\n\n负责部门：体育场馆管理中心\n\n办理网址：https://ehall.neu.edu.cn//db_portal/guide?id=9AA0F0B8-B2DA-4723-9646-A461EFF1D920','怎么预约东北大学的游泳馆？'],
      ['业务名称：游泳馆预约\n\n办理流程：用于游泳馆预约\n\n办理条件：东北大学师生\n\n负责部门：体育场馆管理中心\n\n办理网址：https://ehall.neu.edu.cn//db_portal/guide?id=9AA0F0B8-B2DA-4723-9646-A461EFF1D920','我怎么进入东北大学游泳馆？'],
      ['业务名称：游泳馆预约\n\n办理流程：用于游泳馆预约\n\n办理条件：东北大学师生\n\n负责部门：体育场馆管理中心\n\n办理网址：https://ehall.neu.edu.cn//db_portal/guide?id=9AA0F0B8-B2DA-4723-9646-A461EFF1D920','我想游泳，应该怎么做？'],
      ['业务名称：游泳馆预约\n\n办理流程：用于游泳馆预约\n\n办理条件：东北大学师生\n\n负责部门：体育场馆管理中心\n\n办理网址：https://ehall.neu.edu.cn//db_portal/guide?id=9AA0F0B8-B2DA-4723-9646-A461EFF1D920','我去哪游泳？'],
      ['业务名称：游泳馆预约\n\n办理流程：用于游泳馆预约\n\n办理条件：东北大学师生\n\n负责部门：体育场馆管理中心\n\n办理网址：https://ehall.neu.edu.cn//db_portal/guide?id=9AA0F0B8-B2DA-4723-9646-A461EFF1D920','你知道怎么预约游泳馆嘛'],
      ['业务名称：游泳馆预约\n\n办理流程：用于游泳馆预约\n\n办理条件：东北大学师生\n\n负责部门：体育场馆管理中心\n\n办理网址：https://ehall.neu.edu.cn//db_portal/guide?id=9AA0F0B8-B2DA-4723-9646-A461EFF1D920','你知道我怎么才能游泳不'],
      ['业务名称：游泳馆预约\n\n办理流程：用于游泳馆预约\n\n办理条件：东北大学师生\n\n负责部门：体育场馆管理中心\n\n办理网址：https://ehall.neu.edu.cn//db_portal/guide?id=9AA0F0B8-B2DA-4723-9646-A461EFF1D920','东大游泳馆可以给校外人员预约吗？'],
      ['业务名称：游泳馆预约\n\n办理流程：用于游泳馆预约\n\n办理条件：东北大学师生\n\n负责部门：体育场馆管理中心\n\n办理网址：https://ehall.neu.edu.cn//db_portal/guide?id=9AA0F0B8-B2DA-4723-9646-A461EFF1D920','进入游泳馆有啥要求？']]
for d in data:
    d[1]=(model.embed_documents([d[1]])[0]/maxNorm)[0].to('cpu')
    print(d[1].shape)
# print('data',data)
texts, embeddings = zip(*data)
print(texts,embeddings)
# 真正传入模型需要 [(doc，对应query的embedding),(),()]
db.add_embeddings(data)
print('添加后的数据量：',db.index.ntotal)
db.save_local(index_name)



# query = "介绍下朱靖波"
# docs = db.similarity_search(query,k=10)
# reordering = LongContextReorder()
# reordering_docs = reordering.transform_documents(docs)
# prompt="请根据以下信息回答问题：\n"
# for i in range(3):
#     prompt+=f'{i+1}.'+docs[i].page_content+'\n'
# print(prompt)
